import poser, os, sys, string, re

os.chdir( os.path.dirname(poser.AppLocation()) )

sys.path.append( os.path.join("Runtime","Python","PoseWorks","GlowWorm","Data"))

import gwrRender, gwrFx6

def isRendered(act):
	if act.IsCamera(): return 0
	elif act.IsDeformer(): return 0
	elif act.IsLight(): return 0
	elif act.IsZone(): return 0
	else: return 1

def getAllMaterials(scene, doHair=1):
	materials = []
	acts = scene.Actors()
	for act in acts:
		if act.IsBodyPart(): continue
		elif not isRendered(act): continue
		elif not doHair and act.IsHairProp(): continue
		else:
			for mat in act.Materials(): materials.append(mat)
	for fig in scene.Figures():
		for mat in fig.Materials(): materials.append(mat)
	return materials

class color:
	def __init__(self, r=0, g=0, b=0):
		self.set(r, g, b)
	def set(self, r, g=-1, b=-1):
		if g != -1:
			self.r, self.g, self.b = r, g, b
		else:
			rgb = r.split()
			self.r, self.g, self.b = float(rgb[0]), float(rgb[1]), float(rgb[2])
	def get(self):
		return repr(self.r) + " " + repr(self.g) + " " + repr(self.b)
	def getTuple(self):
		return (self.r, self.g, self.b)

class sceneInfo:
	def __init__(self, scene):
		self.scene = scene
		self.figures = self.scene.Figures()
		self.props = self.getProps()
		self.hair = self.getHairProps()
		self.particles = self.getEffectFigures()
		self.propMats = self.getGroupMaterials(self.props)
		self.hairMats = self.getGroupMaterials(self.hair)
		self.figureMats = self.getGroupMaterials(self.figures)
		self.particleMats = self.getEffectMaterials()
		self.elements = self.figures + self.props + self.hair + self.particles
		self.materials = self.propMats + self.hairMats + self.figureMats + self.particleMats
	def getEffectFigures(self):
		figures = []
		for fig in self.figures:
			if self.effectFilter(fig):
				figures.append(fig)
				self.figures.remove(fig)
		return figures
	def getEffectMaterials(self):
		materials = []
		for fig in self.particles:
			for act in fig.Actors():
				for mat in act.Materials():
					materials.append(mat)
		return materials
	def getProps(self):
		props = []
		for actor in self.scene.Actors():
			if actor.IsProp() and not actor.IsHairProp():
				props.append(actor)
		return props
	def getHairProps(self):
		hairprops = []
		for actor in self.scene.Actors():
			if self.hairFilter(actor):
				hairprops.append(actor)
		return hairprops
	def getGroupMaterials(self, group):
		materials = []
		for member in group:
                        try:
                                for mat in member.Materials():
                                        materials.append(mat)
                        except:
                                continue
                        """
			if not hasattr(member, "Materials"):
				continue
			else:
				member_materials = member.Materials()
				if member_materials and isinstance(member_materials, list):
					for mat in member_materials:
						materials.append(mat)
			"""
		return materials
	def hairFilter(self, actor):
		if actor.IsHairProp() and actor.IsProp():
			return 1
		else:
			return 0
	def effectFilter(self, figure):
		if figure.GeomFileName()[-4:].lower() == 'prt2':
			return 1
		elif figure.ParentActor().Parameter("PRT3_Type"):
                        return 1
		else:
			return 0

class occlusionOptions:
	maxDist = 7.74
	def __init__(self, samples=3.0, maxDist=14, rayBias=0.2064, strength=1.0):
		self.samples = samples
		self.maxDist = maxDist
		self.rayBias = rayBias
		self.strength = strength
	def configure(self, samples=None, maxDist=None, rayBias=None, strength=None):
		if samples: self.samples = samples
		if maxDist: self.maxDist = maxDist
		if rayBias: self.rayBias = rayBias
		if strength: self.strength = strength
	def get(self):
		return vars(self)
	def set(self, *args, **kwds):
		self.configure(kwds)

class renderPass:
	def __init__(self, id="CustomPass", fx6=None, shadingPassShadows=0, passPerLight=0, passPerObject=0, koColor=color(0,0,0), lightingOnly=0, diffuseOnly=0, showHair=0, showParticles=0, hairOnly=0, zbuffer=0, particlesOnly=0, shadows=0, diffuse=0, specular=0, ambient=0, translucence=0, reflection=0, refraction=0, altdiffuse=0, altspecular=0, occlusion=0, shadowOnly=0, aoOptions=occlusionOptions()):
		self.id = id
		self.Kd = diffuse
		self.Ks = specular
		self.Ka = ambient
		self.Ku = translucence
		self.Kr = reflection
		self.Kf = refraction
		self.Kad = altdiffuse
		self.Kas = altspecular
		self.Fx = showParticles
		self.Ha = showHair
		self.shadows = shadows
		self.lighting = lightingOnly
		self.hair = hairOnly
		self.particles = particlesOnly
		self.zbuffer = zbuffer
		self.occlusion = occlusion
		self.shadowOnly = shadowOnly
		self.diffuseOnly = diffuseOnly
		self.aoOptions = aoOptions
		self.koColor = koColor
		self.passPerLight = passPerLight
		self.shadingPassShadows = shadingPassShadows
		self.passPerObject = passPerObject
		self.fx6 = fx6
	def get(self):
		d = dict(vars(self))
		d["koColor"] = self.koColor.get()
		d["aoOptions"] = self.aoOptions.get()
		return d
	def set(self, options):
		for k,v in options.items():
			setattr(self, k, v)
	def setMaterial(self, material):
		st = material.ShaderTree()
		root = st.Node(0)
		if not self.Kd:
			root.InputByInternalName('Diffuse_Value').SetFloat(0)
			root.InputByInternalName('Diffuse_Color').Disconnect()
			root.InputByInternalName('Diffuse_Value').Disconnect()
		if not self.Ks:
			root.InputByInternalName('Highlight_Value').SetFloat(0)
			root.InputByInternalName('Highlight_Color').Disconnect()
			root.InputByInternalName('Highlight_Value').Disconnect()
		if not self.Ka:
			root.InputByInternalName('Ambient_Value').SetFloat(0)
			root.InputByInternalName('Ambient_Color').Disconnect()
			root.InputByInternalName('Ambient_Value').Disconnect()
		if not self.Ku:
			root.InputByInternalName('Translucence_Value').SetFloat(0)
			root.InputByInternalName('Translucence_Color').Disconnect()
			root.InputByInternalName('Translucence_Value').Disconnect()
		if not self.Kr:
			root.InputByInternalName('Reflection_Value').SetFloat(0)
			root.InputByInternalName('Reflection_Color').Disconnect()
			root.InputByInternalName('Reflection_Value').Disconnect()
		if not self.Kf:
			root.InputByInternalName('Refraction_Value').SetFloat(0)
			root.InputByInternalName('Refraction_Color').Disconnect()
			root.InputByInternalName('Refraction_Value').Disconnect()
		if not self.Kad:
			root.InputByInternalName('AlternateDiffuse').SetColor(0,0,0)
			root.InputByInternalName('AlternateDiffuse').Disconnect()
		if not self.Kas:
			root.InputByInternalName('AlternateSpecular').SetColor(0,0,0)
			root.InputByInternalName('AlternateSpecular').Disconnect()
		gwrFx6.removeExcess(material)
	def isNotTransparency(self, input):
		i = input.InternalName()
		if i == 'Transparency_Max': return 0
		elif i == 'Transparency_Min': return 0
		elif i == 'Transparency_Falloff': return 0
		else: return 1
	def clearNonTransparency(self, root):
		for ip in root.Inputs():
			if self.isNotTransparency(ip): ip.Disconnect()
	def knockoutMaterial(self, material, useAmbient=1):
		st = material.ShaderTree()
		root = st.Node(0)
		self.setMaterial(material)
		if useAmbient:
			av = root.InputByInternalName('Ambient_Value')
			av.SetFloat(1)
			av.Disconnect()
			ac = root.InputByInternalName('Ambient_Color')
			ac.SetColor(self.koColor.r, self.koColor.g, self.koColor.b)
			ac.Disconnect()
		else:
			ad = root.InputByInternalName('AlternateDiffuse')
			ad.SetColor(1,1,1)
			sc = st.CreateNode(poser.kNodeTypeCodeSIMPLECOLOR)
			sc.Inputs()[0].SetColor(self.koColor.r, self.koColor.g, self.koColor.b)
			sc.ConnectToInput(ad)
	def setKoColor(self, rgb):
		self.koColor = rgb
	def configureAO(self, *args, **kwds):
		self.aoOptions.configure(kwds)
	def setShadowOnly(self, scene):
		self.doShadows(scene, shadowOnly=1)
	def noShadows(self, scene):
		rs = gwrRender.renderSettings(gwrRender.fireFly())
		rs.forceManual(scene)
		rs.setFlag("shadowOnlyRender",0)
		rs.setFlag("doShadows",0)
		rs.tempFile = "gwrShadow.cm2"
		rs.apply(scene)
		rs.cleanup()
	def doShadows(self, scene, shadowOnly=0):
		rs = gwrRender.renderSettings(gwrRender.fireFly())
		rs.forceManual(scene)
		rs.setFlag("doShadows",1)
		rs.setFlag("shadowOnlyRender",shadowOnly)
		rs.tempFile = "gwrShadow.cm2"
		rs.apply(scene)
		rs.cleanup()
	def setLightingOnly(self, scene):
		if not self.shadingPassShadows:
			self.noShadows(scene)
		else:
			self.doShadows(scene)
		self.setNoFX(scene)
		if self.passPerLight:
			self.forceBGColor(scene, 0, 0, 0)
		else:
			self.forceBGColor(scene, 0.5, 0.5, 0.5)
		for mat in self.info.materials:
			st = mat.ShaderTree()
			root = st.Node(0)
			self.Kd = 1
			dc = root.InputByInternalName('Diffuse_Color')
			dc.SetColor(0.5,0.5,0.5)
			dc.Disconnect()
			self.setMaterial(mat)
	def setDIF(self, scene):
		for light in scene.Lights(): light.Delete()
		self.setNoFX(scene)
		self.Ka = 1
		self.Kd = 0
		for mat in self.info.materials:
			st = mat.ShaderTree()
			root = st.Node(0)
			dc = root.InputByInternalName('Diffuse_Color')
			dcol = dc.Value()
			dcIn = dc.InNode()
			dv = root.InputByInternalName('Diffuse_Value')
			dval = dv.Value()
			dv.SetFloat(0)
			root.InputByInternalName('Ambient_Value').SetFloat(dval)
			ac = root.InputByInternalName('Ambient_Color')
			ac.SetColor(dcol[0],dcol[1],dcol[2])
			if dcIn: dcIn.ConnectToInput(ac)
			self.Ka = 1
			self.setMaterial(mat)
	def setAO(self, scene):
		for light in scene.Lights(): light.Delete()
		rs = gwrRender.renderSettings(gwrRender.fireFly())
		rs.forceManual(scene)
		rs.setFlag("allowRayTracing",1)
		rs.setFlag("doShadows",0)
		rs.tempFile = "gwrOcclusion.cm2"
		rs.apply(scene)
		rs.cleanup()
		self.forceBGColor(scene,1,1,1)
		self.setNoFX(scene)
		for mat in (self.info.hairMats):
			self.blankMat(mat)
		for material in (self.info.figureMats + self.info.propMats):
			self.knockoutMaterial(material, useAmbient=1)
			st = material.ShaderTree()
			root = st.Node(0)
			ad = root.InputByInternalName('Ambient_Color')
			ad.SetColor(1,1,1)
			ao = st.CreateNode(poser.kNodeTypeCodeAMBIENTOCCLUSION)
			ao.InputByInternalName('Samples').SetFloat(self.aoOptions.samples)
			ao.InputByInternalName('MaxDist').SetFloat(self.aoOptions.maxDist)
			ao.InputByInternalName('RayBias').SetFloat(self.aoOptions.rayBias)
			ao.InputByInternalName('Strength').SetFloat(self.aoOptions.strength)
			ao.ConnectToInput(ad)
	def setZB(self, scene, materials):
		for light in scene.Lights(): light.Delete()
		import gwoClip
		rs = gwrRender.renderSettings(gwrRender.fireFly())
		rs.forceManual(scene)
		rs.tempFile = "gwrZbuffer.cm2"
		rs.apply(scene)
		rs.cleanup()
		self.forceBGColor(scene,1,1,1)
		self.setNoFX(scene)
		st = scene.AtmosphereShaderTree()
		atmos = st.Node(0)
		atmos.InputByInternalName("DepthCue_On").SetFloat(1)
		atmos.InputByInternalName("DepthCue_Color").SetColor(1,1,1)
		gwoClip.depthCue(scene)
		for material in materials:
			self.knockoutMaterial(material)
	def setFX(self, scene):
		self.noShadows(scene)
		for act in (self.info.hair + self.info.props): act.SetVisibleInReflections(0)
		for fig in self.info.figures:
			for act in fig.Actors(): act.SetVisibleInReflections(0)
		for mat in (self.info.hairMats + self.info.figureMats + self.info.propMats):
			self.knockoutMaterial(mat)
		self.forceBGColor(scene, self.koColor.r, self.koColor.g, self.koColor.b)
	def setNoFX(self, scene):
		self.showParticles(scene, 0)
	def setHair(self, scene):
		self.noShadows(scene)
		self.setNoFX(scene)
		for act in self.info.props: act.SetVisibleInReflections(0)
		for fig in (self.info.figures + self.info.particles):
			for act in fig.Actors(): act.SetVisibleInReflections(0)
		for mat in (self.info.particleMats + self.info.figureMats + self.info.propMats):
			self.knockoutMaterial(mat)
		self.forceBGColor(scene, self.koColor.r, self.koColor.g, self.koColor.b)
	def setNoHair(self, scene):
		self.showHair(scene, 0)
	def showParticles(self, scene, bool):
		for effect in self.info.particles:
			if effect.Visible(): effect.SetVisible(bool)
	def showHair(self, scene, bool):
		for hairProp in self.info.hair:
			if hairProp.Visible(): hairProp.SetVisible(bool)
	def blankMat(self, material):
		st = material.ShaderTree()
		root = st.Node(0)
		root.InputByInternalName('Ambient_Color').SetColor(1,1,1)
		root.InputByInternalName('Ambient_Value').SetFloat(1.0)
	def forceBGColor(self, scene, r, g, b):
		bgSt = scene.BackgroundShaderTree()
		root = bgSt.Node(0)
		try:
			bgCol = bgSt.NodeByInternalName("BG Color")
		except:
			bgCol = bgSt.CreateNode(poser.kNodeTypeCodeSIMPLECOLOR)
		bgCol.Input(0).Disconnect()
		bgCol.Input(0).SetColor(r,g,b)
		root.Input(1).Disconnect()
		root.Input(1).SetColor(0,0,0)
		root.Input(2).Disconnect()
		root.Input(2).SetColor(0,0,0)
		scene.SetBackgroundColor(r,g,b)
		scene.SetRenderOnBGColor(1)
	def apply(self, scene):
		self.info = sceneInfo(scene)
		if self.shadowOnly:
			self.setShadowOnly( scene )
		elif self.zbuffer:
			self.setZB( scene, self.info.materials )
		elif self.particles:
			self.setFX( scene )
		elif self.hair:
			self.setHair( scene )
		elif self.diffuseOnly:
			self.setDIF( scene )
		elif self.lighting:
			self.setLightingOnly( scene )
		elif self.occlusion:
			self.setAO(scene)
		else:
			for material in (self.info.propMats + self.info.hairMats + self.info.figureMats):
				self.setMaterial(material)
			if self.shadows:
				self.doShadows(scene)
			else:
				self.noShadows(scene)
			if not self.Kd:
				self.forceBGColor(scene, 0, 0, 0)
				#if not self.Kad:
				#	for light in scene.Lights(): light.Delete()
			if not self.Fx:
				self.setNoFX(scene)
			if not self.Ha:
				self.setNoHair(scene)
		if self.fx6 and self.fx6 != "None":
			gwrFx6.LoadFX6OnScene(self.fx6, scene)

hairPass = renderPass(id="fur", hairOnly=1)
effectsPass = renderPass(id="fx", particlesOnly=1)
depthPass = renderPass(id="zb", zbuffer=1)
shadowPass = renderPass(id="shd", shadowOnly=1)
occlusionPass = renderPass(id="ao", occlusion=1)
diffusePass = renderPass(id="dif", diffuseOnly=1)
specularPass = renderPass(id="spec", showHair=1, specular=1, shadows=0)
ambientPass = renderPass(id="amb", showHair=1, ambient=1)
glowPass = renderPass(id="glow", showHair=1, ambient=1, translucence=1)
beautyPass = renderPass(id="beau", showHair=1, diffuse=1)
lightingPass = renderPass(id="lite", lightingOnly=1)
normalPass = renderPass(id="nr", showHair=1, showParticles=1, shadows=1, diffuse=1, specular=1, ambient=1, translucence=1, reflection=1, refraction=1)

def openAndRender(filePath, renderPass, frame="USEDOC"):
	poser.CloseDocument(1)
	poser.OpenDocument(filePath)
	scene = poser.Scene()
	if frame != "USEDOC": scene.SetFrame(frame)
	renderPass.apply(scene)
	scene.Render()

def singleLightScene(scene, light):
	iN = light.InternalName()
	fi = open("temp.lt2", 'w')
	fi.write("""{\nversion\n\t{\n\tnumber 6\n\t}\nlight %s\n\t{\n\tlightOn 1\n\t}\n}""" % iN )
	fi.close()
	scene.LoadLibraryLight("temp.lt2")
	os.remove("temp.lt2")

if __name__ == '__main__':
	pz3Path = "C:\\Documents and Settings\\Dan\\Desktop\\rpTest.pz3"
	scene = poser.Scene()
	#normalPass.apply(scene)
	#occlusionPass.apply(scene)
	#scene.Render()
	#openAndRender(pz3Path, zbPass)

pass
